#include "LinearSystemSolver.h"

LinearSystemSolver::LinearSystemSolver(void)
{
}

LinearSystemSolver::LinearSystemSolver( ParticleSystem *ps, ConstraintSystem *cs )
{
	SystemInitialize(ps,cs);
}

LinearSystemSolver::~LinearSystemSolver(void)
{
	Clear();
}

void LinearSystemSolver::Clear()
{
	for (int i=0;i<scale;i++)
	{
		delete [] A[i];
		delete [] GS[i];
	}

	for (int i=0;i<cNum;i++)
	{
		delete [] CGrad[i];
	}

	delete [] A;
	delete [] CGrad;
	delete [] b;
	delete [] HO;
	delete [] cache;
	delete [] output;
	delete [] residual;
	delete [] residual_cache;
	delete [] p;
	delete [] GS;

#ifdef IMKL
	delete [] svd.A;
	delete [] svd.superb;
	delete [] svd.u;
	delete [] svd.v;
	delete [] svd.s;

	for (int i=0;i<scale;i++)
	{
		delete [] svd.cache[i];
		delete [] svd.pseInv[i];
	}

	delete [] svd.cache;
	delete [] svd.pseInv;
#endif
}

void LinearSystemSolver::SystemInitialize( ParticleSystem *ps, ConstraintSystem *cs )
{
	this->ps=ps;
	this->cs=cs;
	scale=ps->GetScale()+cs->GetScale();
	pNum=ps->GetScale();
	cNum=cs->GetScale();

	A=(double**)malloc(scale*sizeof(double*));
	GS=(double**)malloc(scale*sizeof(double*));
	CGrad=(double**)malloc(cNum*sizeof(double*));
	b=(double*)malloc(scale*sizeof(double));
	cache=(double*)malloc(scale*sizeof(double));
	HO=(double*)malloc(cNum*sizeof(double));
	output=(double*)malloc(scale*sizeof(double));
	residual=(double*)malloc(scale*sizeof(double));
	residual_cache=(double*)malloc(scale*sizeof(double));
	p=(double*)malloc(scale*sizeof(double));

#ifdef IMKL
	svd.m=svd.n=scale;
	svd.lda=svd.ldu=svd.ldvt=scale;
	svd.A=(double*)malloc(scale*scale*sizeof(double));
	svd.u=(double*)malloc(scale*scale*sizeof(double));
	svd.v=(double*)malloc(scale*scale*sizeof(double));
	svd.superb=(double*)malloc(scale*scale*sizeof(double));
	svd.s=(double*)malloc(scale*sizeof(double));
	svd.cache=(double**)malloc(scale*sizeof(double*));
	svd.pseInv=(double**)malloc(scale*sizeof(double*));

	for (int i=0;i<scale;i++)
	{
		svd.cache[i]=(double*)malloc(scale*sizeof(double));
		svd.pseInv[i]=(double*)malloc(scale*sizeof(double));
	}
#endif

	for (int i=0;i<cNum;i++)
	{
		CGrad[i]=(double*)malloc(pNum*sizeof(double));
		HO[i]=0;

		for (int j=0;j<pNum;j++)
		{
			CGrad[i][j]=0;
		}
	}

	for (int i=0;i<scale;i++)
	{
		A[i]=(double*)malloc(scale*sizeof(double));
		GS[i]=(double*)malloc((scale+1)*sizeof(double));
		b[i]=0;
		cache[i]=0;
		output[i]=0;
		residual[i]=0;
		residual_cache[i]=0;
		p[i]=0;

		for (int j=0;j<scale;j++)
		{
			A[i][j]=0;
		}
	}
}

void LinearSystemSolver::SystemConstruction()
{
	MassBlockConstruction();
	cs->GetConstraintJacobian(CGrad,&(ps->particles));
	CGradBlockConstruction();
	CGradTransposeBlockConstruction();
	BVectorConstruction();	
}

void LinearSystemSolver::MassBlockConstruction()
{
	int index=0;

	for (int i=0;i<ps->particles.size();i++)
	{
		A[index][index]=ps->particles[i]->mass;
		index++;
		A[index][index]=ps->particles[i]->mass;
		index++;
	}
}

void LinearSystemSolver::CGradBlockConstruction()
{
	int index=pNum;

	for (int i=0;i<cNum;i++)
	{
		for (int j=0;j<pNum;j++)
		{
			A[index+i][j]=CGrad[i][j];
		}
	}
}

void LinearSystemSolver::CGradTransposeBlockConstruction()
{
	int index=pNum;

	for (int i=0;i<pNum;i++)
	{
		for (int j=0;j<cNum;j++)
		{
			A[i][j+index]=CGrad[j][i];
		}
	}
	
}

void LinearSystemSolver::ExtForceConstruction()
{
	int index=0;

	for (int i=0;i<ps->particles.size();i++)
	{
		b[index]=0;
		index++;
		b[index]=9.8*ps->particles[i]->mass;
		index++;
	}
}

void LinearSystemSolver::HarmonicOscilatorConstruction()
{
	cs->GetHarmonicOsciliator(HO,&(ps->particles));

	int index=pNum;

	for (int i=0;i<cNum;i++)
	{
		b[index+i]=HO[i];
	}
}

void LinearSystemSolver::BVectorConstruction()
{
	ExtForceConstruction();
	HarmonicOscilatorConstruction();
}

void LinearSystemSolver::Solve()
{
	SystemConstruction();
	SystemIterativeSolve();
}

void LinearSystemSolver::GetAccleration( double *accleration )
{
	for (int i=0;i<pNum;i++)
	{
		accleration[i]=output[i];
	}
}

void LinearSystemSolver::SystemIterativeSolve()
{
	switch (solverType)
	{
	case CONJUGATEGRADIENT:
		ConjugateGradIter();
		break;

	case GAUSSJORDAN:
		GaussJordan();
		break;

	case SINGLEVALUEDECOMPOSITION:
		SVDSolver();
		break;
	}
	
}

void LinearSystemSolver::JacobiIter()
{
	int loop;

	for (loop=0;loop<1000;loop++)
	{
		for (int i=0;i<scale;i++)
		{
			double sum=0;

			for (int j=0;j<scale;j++)
			{
				if (i!=j)
				{
					sum+=A[i][j]*output[j];
				}
			}

			if (A[i][i]!=0)
			{
				cache[i]=(b[i]-sum)/A[i][i];
			}
		}

		double r=abs(cache[0]-output[0]);

		for (int i=0;i<scale;i++)
		{
			if (r<abs(cache[i]-output[i]))
			{
				r=abs(cache[i]-output[i]);
			}
			output[i]=cache[i];
		}

		if (r<1e-5)
		{
			break;
		}
	}

	printf("%d\n",loop);
}

void LinearSystemSolver::ConjugateGradIter()
{
	for (int i=0;i<scale;i++)
	{
		double sum=0;

		for (int j=0;j<scale;j++)
		{
			sum+=A[i][j]*output[j];
		}

		residual[i]=b[i]-sum;
		p[i]=residual[i];
	}

	int loop;

	for (loop=0;loop<100;loop++)
	{
		double alpha, sumU=0, sumD=0;

		for (int i=0;i<scale;i++)
		{
			sumU+=residual[i]*residual[i];
			double sumT=0;

			for (int j=0;j<scale;j++)
			{
				sumT+=p[j]*A[j][i];
			}

			sumD+=sumT*p[i];
		}

		alpha=sumU/sumD;

		for (int i=0;i<scale;i++)
		{
			cache[i]=output[i]+alpha*p[i];
			
			double sum=0;

			for (int j=0;j<scale;j++)
			{
				sum+=A[i][j]*p[j];
			}

			residual_cache[i]=residual[i]-alpha*sum;
		}

		/************************************************************************/
		/* check here                                                                     */
		/************************************************************************/

		double r=abs(residual_cache[0]);

		for (int i=0;i<scale;i++)
		{
			if (r<abs(residual_cache[i]))
			{
				r=abs(residual_cache[i]);
			}
		}

		if (r<1e-5)
		{
			double *swap;
			swap=output;
			output=cache;
			cache=swap;
			return;
		}

		double beta;
		sumU=0;
		sumD=0;

		for (int i=0;i<scale;i++)
		{
			sumU+=residual_cache[i]*residual_cache[i];
			sumD+=residual[i]*residual[i];
		}

		beta=sumU/sumD;

		for (int i=0;i<scale;i++)
		{
			p[i]=residual_cache[i]+beta*p[i];
		}

		double *swap;
		swap=residual;
		residual=residual_cache;
		residual_cache=swap;

		swap=output;
		output=cache;
		cache=swap;
	}
}

void LinearSystemSolver::GaussJordan()
{

	{//Matrix Extension Construction
		for (int i=0;i<scale;i++)
		{
			for (int j=0;j<scale;j++)
			{
				GS[i][j]=A[i][j];
			}

			GS[i][scale]=b[i];
		}
	}

	{//Forward Substitution
		int max;
		double swap;

		for (int i=0;i<scale;i++)
		{
			max=i;

			for (int j=i+1;j<scale;j++)
			{
				if (abs(GS[j][i])>abs(GS[max][i]))
				{
					max=j;
				}
			}

			for (int j=0;j<=scale;j++)
			{
				swap=GS[max][j];
				GS[max][j]=GS[i][j];
				GS[i][j]=swap;
			}

			if (abs(GS[i][i]) < 1e-5)
			{
				printf("Singular\n");
			}

			for (int j=scale;j>=i;j--)
			{
				for (int k=i+1;k<scale;k++)
				{
					GS[k][j]-=GS[k][i]/GS[i][i]*GS[i][j];
				}
			}
		}
	}

	{//Reverse Elimination
		for (int i=scale-1;i>=0;i--)
		{
			GS[i][scale]=GS[i][scale]/GS[i][i];
			GS[i][i]=1;

			for (int j=i-1;j>=0;j--)
			{
				GS[j][scale]-=GS[j][i]*GS[i][scale];
				GS[j][i]=0;
			}
		}
	}

	{
		for (int i=0;i<scale;i++)
		{
			output[i]=GS[i][scale];
		}	
	}

}

void LinearSystemSolver::SVD()
{

#ifdef IMKL
	for (int i=0;i<scale;i++)
	{
		for (int j=0;j<scale;j++)
		{
			svd.A[I(i,j)]=A[i][j];
		}
	}

	LAPACKE_dgesvd( LAPACK_ROW_MAJOR, 'A', 'A', svd.m, svd.n, svd.A, svd.lda, svd.s, svd.u, svd.ldu, svd.v, svd.ldvt, svd.superb );
#endif
	
}

int LinearSystemSolver::I( int i, int j )
{
	return i*scale+j;
}

void LinearSystemSolver::SVDSolver()
{
	SVD();
	PseudoInverse();
}

void LinearSystemSolver::PseudoInverse()
{
#ifdef IMKL
	for (int i=0;i<scale;i++)
	{
		if (svd.s[i]>1e-5)
		{
			svd.s[i]=1.0/svd.s[i];
		}else{
			break;
		}
	}

	for (int i=0;i<scale;i++)
	{
		for (int j=0;j<scale;j++)
		{
			svd.cache[i][j]=svd.u[I(i,j)]*svd.s[j];
		}
	}

	for (int i=0;i<scale;i++)
	{
		for (int j=0;j<scale;j++)
		{
			double sum=0;
			for (int k=0;k<scale;k++)
			{
				sum+=svd.cache[i][k]*svd.v[I(k,j)];
			}
			svd.pseInv[i][j]=sum;
		}
	}

	for (int i=0;i<scale;i++)
	{
		double sum=0;
		for (int j=0;j<scale;j++)
		{
			sum+=svd.pseInv[i][j]*b[j];
		}
		output[i]=sum;
	}

#endif
}
